from nl2sql.model.llm import LLM
from nl2sql.prompt.generate_sql import sql_prompt


class SqlGeneration:
    def __init__(self, model: LLM) -> None:
        self.prompt = sql_prompt
        self.llm = model

    def generate_sql(self, query: str, data_info: str) -> str:
        """
        用于拆解用户的问题，为子问题
        :param data_info:
        :param query:
        :return: None or List of String
        """
        # 把Prompt模板，进行填充
        filled_prompt = self.prompt.format(query=query, data_info=data_info)
        # 让大模型对填充好的模板，进行推理
        print('DEBUG', filled_prompt)
        result = self.llm.chat(filled_prompt)
        return result
